import numpy as np
import torch
from torch.utils.data import Dataset

from utils import text_to_sequence


class InvHashDataset(Dataset):
    """
    A PyTorch Dataset class to be used in a PyTorch DataLoader to create batches.
    """

    def __init__(self, split):
        filename = '{}.txt'.format(split)
        with open(filename, 'r') as file:
            data = file.readlines()

        self.samples = data

    def __getitem__(self, i):
        sample = self.samples[i]
        token = sample.strip().split()
        x = torch.FloatTensor(bytearray.fromhex(token[1]))
        y = torch.FloatTensor(text_to_sequence(token[0]))
        return x, y

    def __len__(self):
        return len(self.samples)

    def shuffle(self):
        np.random.shuffle(self.samples)


if __name__ == "__main__":
    data = InvHashDataset('train')
    print(data[0])
